from typing import Dict, List, Optional, Any
from loguru import logger
from pathlib import Path
import json
import yaml
from datetime import datetime, timezone
import os
import logging
import sqlite3
import requests

from .osv_searcher import OSVSearcher
from .osv_parser import OsvParser

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class OSVCollector:
    """OSV数据收集器"""
    
    def __init__(self, db_path: str = "vulnerabilities.db"):
        """初始化收集器"""
        self.db_path = db_path
        self.vulnerabilities = []
        self.reviewed_vulnerabilities = []
        self.init_db()
        self.load_vulnerabilities()
        self.load_reviewed_vulnerabilities()
    
    def init_db(self):
        """初始化数据库"""
        conn = sqlite3.connect(self.db_path)
        c = conn.cursor()
        c.execute('''
            CREATE TABLE IF NOT EXISTS vulnerabilities
            (id TEXT PRIMARY KEY,
             package TEXT,
             severity TEXT,
             published TEXT,
             description TEXT,
             source TEXT)
        ''')
        conn.commit()
        conn.close()
    
    def load_vulnerabilities(self):
        """加载漏洞数据"""
        try:
            conn = sqlite3.connect(self.db_path)
            c = conn.cursor()
            c.execute('SELECT * FROM vulnerabilities')
            rows = c.fetchall()
            self.vulnerabilities = [
                {
                    'id': row[0],
                    'package': row[1],
                    'severity': row[2],
                    'published': row[3],
                    'description': row[4],
                    'source': row[5]
                }
                for row in rows
            ]
            conn.close()
            logger.info(f"Loaded {len(self.vulnerabilities)} vulnerabilities from database")
        except Exception as e:
            logger.error(f"Error loading vulnerabilities: {e}")
            self.vulnerabilities = []
    
    def load_reviewed_vulnerabilities(self):
        """加载已审核的漏洞数据"""
        try:
            conn = sqlite3.connect(self.db_path)
            c = conn.cursor()
            c.execute('SELECT * FROM reviewed_vulnerabilities')
            rows = c.fetchall()
            self.reviewed_vulnerabilities = [
                {
                    'id': row[0],
                    'package': row[1],
                    'severity': row[2],
                    'published': row[3],
                    'description': row[4],
                    'source': row[5]
                }
                for row in rows
            ]
            conn.close()
            reviewed_count = len(self.reviewed_vulnerabilities)
            logger.info(f"Total reviewed vulnerabilities loaded: {reviewed_count}")
        except Exception as e:
            logger.error(f"Error loading reviewed vulnerabilities: {e}")
            self.reviewed_vulnerabilities = []
    
    def get_stats(self):
        """获取漏洞统计信息"""
        try:
            all_vulns = self.vulnerabilities
            
            total = len(all_vulns)
            
            by_severity = {
                'high': 0,
                'medium': 0,
                'low': 0,
                'critical': 0,
                'fixed': 0
            }
            
            for vuln in all_vulns:
                severity = vuln.get('severity', 'unknown')
                
                if isinstance(severity, list):
                    severity = severity[0] if severity else 'unknown'
                severity = str(severity).lower()
                
                is_fixed = vuln.get('fixed', False)
                
                if is_fixed:
                    by_severity['fixed'] += 1
                elif severity in by_severity:
                    by_severity[severity] += 1
                else:
                    by_severity['low'] += 1
            
            logger.info(f"Generated stats: total={total}, by_severity={by_severity}")
            
            return {
                'total_vulnerabilities': total,
                'by_severity': by_severity
            }
        except Exception as e:
            logger.error(f"获取统计信息失败: {str(e)}")
            return {
                'total_vulnerabilities': 0,
                'by_severity': {
                    'high': 0,
                    'medium': 0,
                    'low': 0,
                    'critical': 0,
                    'fixed': 0
                }
            }
    
    def get_recent_vulnerabilities(self, limit: int = 10, page: int = 1) -> List[Dict[str, Any]]:
        """获取最近的漏洞（带分页）"""
        try:
            start_idx = (page - 1) * limit
            end_idx = start_idx + limit
            return self.vulnerabilities[start_idx:end_idx]
        except Exception as e:
            logger.error(f"Error getting recent vulnerabilities: {e}")
            return []
    
    def get_total_vulnerabilities(self) -> int:
        """获取漏洞总数"""
        return len(self.vulnerabilities)
    
    def get_vulnerability_by_id(self, vuln_id: str) -> Optional[Dict[str, Any]]:
        """获取特定漏洞"""
        try:
            for vuln in self.vulnerabilities:
                if vuln.get('id') == vuln_id:
                    return vuln
            return None
        except Exception as e:
            logger.error(f"Error getting vulnerability by ID: {e}")
            return None
    
    def search_vulnerabilities(self, query: str) -> List[Dict[str, Any]]:
        """搜索漏洞"""
        try:
            query = query.lower()
            results = []
            for vuln in self.vulnerabilities:
                if (query in vuln.get('id', '').lower() or
                    query in vuln.get('package', '').lower() or
                    query in vuln.get('description', '').lower()):
                    results.append(vuln)
            return results
        except Exception as e:
            logger.error(f"Error searching vulnerabilities: {e}")
            return []
    
    def get_samples(self, vuln_id: str) -> List[Dict[str, Any]]:
        """获取样本数据"""
        try:
            samples = []
            for vuln in self.vulnerabilities:
                if vuln.get('id') == vuln_id:
                    if 'samples' in vuln:
                        samples.extend(vuln['samples'])
            return samples
        except Exception as e:
            logger.error(f"Error getting samples: {e}")
            return []

    def _load_config(self):
        """加载配置"""
        try:
            config_path = Path("config.yml")
            if config_path.exists():
                with open(config_path, 'r', encoding='utf-8') as f:
                    self.config = yaml.safe_load(f)
            else:
                logger.warning("配置文件不存在，使用默认配置")
                self.config = {
                    'storage': {
                        'data_dir': 'data',
                        'temp_dir': 'temp',
                        'cache_dir': 'cache'
                    }
                }
        except Exception as e:
            logger.error(f"加载配置失败: {e}")
            self.config = {
                'storage': {
                    'data_dir': 'data',
                    'temp_dir': 'temp',
                    'cache_dir': 'cache'
                }
            }
    
    def _init_data(self):
        """初始化数据"""
        try:
            self._load_existing_data()
            self._collect_new_data()
        except Exception as e:
            logger.error(f"初始化数据失败: {e}")
    
    def _load_existing_data(self):
        """加载现有数据"""
        data_dir = Path(self.config['storage']['data_dir'])
        if not data_dir.exists():
            data_dir.mkdir(parents=True)
            self.searcher = OSVSearcher(vulnerabilities=[], packages=[])
            return
        
        # 加载漏洞数据
        vuln_file = data_dir / 'vulnerabilities.json'
        vulnerabilities = []
        if vuln_file.exists():
            with open(vuln_file, 'r', encoding='utf-8') as f:
                vulnerabilities = json.load(f)
        
        # 加载包数据
        package_file = data_dir / 'packages.json'
        packages = []
        if package_file.exists():
            with open(package_file, 'r', encoding='utf-8') as f:
                packages = json.load(f)
        
        self.searcher = OSVSearcher(vulnerabilities=vulnerabilities, packages=packages)
        
        # 初始化解析器
        osv_file = data_dir / 'osv.json'
        if osv_file.exists():
            with open(osv_file, 'r', encoding='utf-8') as f:
                osv_data = json.load(f)
            self.parser = OsvParser(osv_json=osv_data)
        else:
            self.parser = OsvParser(osv_json={})
    
    def _collect_new_data(self):
        """收集新数据"""
        try:
            # 这里应该实现实际的数据收集逻辑
            # 目前使用模拟数据
            self.searcher.vulnerabilities = [
                {
                    "id": "CVE-2024-1234",
                    "severity": "high",
                    "package": "package-name@1.0.0",
                    "description": "A critical vulnerability...",
                    "published": "2024-03-31"
                }
            ]
            self.searcher.packages = [
                {
                    "name": "package-name",
                    "version": "1.0.0",
                    "ecosystem": "npm"
                }
            ]
            
            # 保存数据
            self._save_data()
        except Exception as e:
            logger.error(f"收集新数据失败: {e}")
    
    def _save_data(self):
        """保存数据"""
        try:
            data_dir = Path(self.config['storage']['data_dir'])
            data_dir.mkdir(parents=True, exist_ok=True)
            
            # 保存漏洞数据
            with open(data_dir / 'vulnerabilities.json', 'w', encoding='utf-8') as f:
                json.dump(self.searcher.vulnerabilities, f, ensure_ascii=False, indent=2)
            
            # 保存包数据
            with open(data_dir / 'packages.json', 'w', encoding='utf-8') as f:
                json.dump(self.searcher.packages, f, ensure_ascii=False, indent=2)
        except Exception as e:
            logger.error(f"保存数据失败: {e}")
    
    def search(self, field: str, value: str) -> List[Dict]:
        """搜索数据"""
        return self.searcher.search(field, value)
    
    def get_statistics(self) -> Dict:
        """获取统计信息"""
        return {
            "total": len(self.searcher.vulnerabilities),
            "high": len([v for v in self.searcher.vulnerabilities if v.get('severity') == 'high']),
            "fixed": len([v for v in self.searcher.vulnerabilities if v.get('fixed')])
        }

    def get_all_vulnerabilities(self) -> List[Dict[str, Any]]:
        """获取所有漏洞"""
        return self.vulnerabilities

    def get_samples(self) -> List[Dict[str, Any]]:
        """获取样本数据"""
        # 这里可以根据实际需求实现样本数据的获取
        return [] 